import "babel-polyfill";
import {binarySearch, islice} from "../../algorithm/array";
import {every, some} from "../../algorithm/iterator";
import Integer, {range} from "../number/integer";
import {fastExponentation} from "../generic-algorithm";
import {gcd} from "./gcd";
import {jacobi} from "./arithmetic";
import {sqrt} from "./nth-root";

const {zero, one, eq, compare, negate, add, subtract, multiply, square, mod, pow, shl} = Integer;

export var small_primes = [
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97
].map(p => Integer[p]);

var prime_bound = Integer[100];

export var initPrimes = function(bound) {
    bound = Integer.from(bound);
    if (compare(prime_bound, bound) >= 0) {
        return;
    }
    if (compare(square(prime_bound), bound) < 0) {
        initPrimes(sqrt(bound));
    }
    var length = small_primes.length;
    var start = add(prime_bound, (prime_bound.lastLimb & 1) + 1);
    for (let n of range(start, add(bound, 1), 2)) {
        if (small_primes::islice(1, length)::every(p => mod(n, p).sign !== 0)) {
            small_primes.push(n);
        }
    }
    prime_bound = bound;
};

initPrimes(10000);

var millerRabin = function(n, base_list = [2]) {
    if (base_list.some(base => !eq(gcd(base, n), one))) {
        return false;
    }
    var m = subtract(n, 1);
    var [s, d] = m.extractEven();
    for (let base of base_list) {
        let x = pow(base, d, n);
        if (eq(x, one) || eq(x, m)) {
            continue;
        }
        let times = 0;
        while (++times < s) {
            x = mod(square(x), n);
            if (eq(x, one)) {
                return false;
            } else if (eq(x, m)) {
                break;
            }
        }
        if (times === s) {
            return false;
        }
    }
    return true;
};

var lucas = function(n) {
    var q = one;
    while (true) {
        q = q.sign > 0 ? negate(q) : subtract(one, q);
        if (jacobi(subtract(one, shl(q, 2)), n) === -1) {
            break;
        }
    }
    return fastExponentation([one, negate(q), one, zero], n, {
        multiply: (x, y) => [
            add(multiply(x[0], y[0]), multiply(x[1], y[2])),
            add(multiply(x[0], y[1]), multiply(x[1], y[3])),
            add(multiply(x[2], y[0]), multiply(x[3], y[2])),
            add(multiply(x[2], y[1]), multiply(x[3], y[3]))
        ].map(x => mod(x, n)),
        square: function(x) {
            var x12 = multiply(x[1], x[2]);
            return [
                add(square(x[0]), x12),
                add(multiply(x[0], x[1]), multiply(x[1], x[3])),
                add(multiply(x[2], x[0]), multiply(x[3], x[2])),
                add(x12, square(x[3]))
            ].map(x => mod(x, n));
        },
        one: [one, zero, zero, one]
    })[0].sign === 0;
};

export var isProbablyPrime = function(n) {
    n = Integer.from(n);
    if (n.sign <= 0) {
        return false;
    } else if (compare(n, prime_bound) <= 0) {
        return small_primes::binarySearch(n, compare) >= 0;
    } else {
        if (small_primes::some(p => mod(n, p).sign === 0)) {
            return false;
        } else if (compare(n, square(prime_bound)) <= 0) {
            return true;
        } else {
            return millerRabin(n) && lucas(n);
        }
    }
};

export var isPrime = isProbablyPrime;

export var generate = function*(bound) {
    if (arguments.length) {
        bound = Integer.from(bound);
    }
    for (let p of small_primes) {
        if (bound && compare(p, bound) > 0) {
            return;
        }
        yield p;
    }
    for (let p of range(add(small_primes[small_primes.length - 1], 2), undefined, 2)) {
        if (bound && compare(p, bound) > 0) {
            return;
        } else if (isPrime(p)) {
            yield p;
        }
    }
};
